refactor: pair exclude_types as canonical NeighborGraph transform; dpa1 graph path supports exclude_types (decision #18)#5733
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds pair-exclusion support across neighbor-list and neighbor-graph paths, extends DPA1 graph-native attention and export/tracing behavior, updates spin routing, and applies matching pair-exclusion metadata in the C++ pt_expt inference seam with new tests. ChangesPython pair exclusion and graph-native DPA1 attention
C++ pt_expt pair-exclusion ingestion seam
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant DescrptDPA1
participant DescrptBlockSeAtten
participant center_edge_pairs
participant segment_softmax
DescrptDPA1->>DescrptBlockSeAtten: call_graph(graph, atype, static_nnei)
DescrptBlockSeAtten->>DescrptBlockSeAtten: apply_pair_exclusion(graph, atype, pair_excl)
DescrptBlockSeAtten->>center_edge_pairs: center_edge_pairs(dst, edge_mask, static_nnei)
DescrptBlockSeAtten->>segment_softmax: segment_softmax(scores, query_edge, mask)
segment_softmax-->>DescrptBlockSeAtten: normalized attention weights
DescrptBlockSeAtten-->>DescrptDPA1: grrg, rot_mat
sequenceDiagram
participant DeepPotPTExptInit
participant DeepPotPTExptCompute
participant buildPairExcludeTable
participant applyPairExclusion
participant applyPairExclusionNlist
DeepPotPTExptInit->>buildPairExcludeTable: pair_exclude_types
DeepPotPTExptCompute->>applyPairExclusion: graph edge_index, edge_mask, atype
DeepPotPTExptCompute->>applyPairExclusionNlist: nlist, atype_ext
applyPairExclusion-->>DeepPotPTExptCompute: filtered edge_mask
applyPairExclusionNlist-->>DeepPotPTExptCompute: filtered nlist
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/tests/pt_expt/descriptor/test_dpa1.py (1)
117-147: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPass
mappingtotorch.export.exporthere
test_exportablestill goes through the dense fallback becauseTestCaseSingleFrameWithNlistsetsnloc=3andnall=4, while this export call passes nomapping. That means the newexclude_typescase only covers the legacy dense exclusion mask, not the graph-nativeapply_pair_exclusionpath. Addmappingto the exported inputs so the parametrization exercises the intended route.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt_expt/descriptor/test_dpa1.py` around lines 117 - 147, test_exportable is missing the graph mapping input, so it still exercises the dense fallback instead of the graph-native exclusion path. Update the export setup in test_exportable to pass mapping into torch.export.export alongside dd0 and the existing inputs, using the test fixture’s mapping source so the exclude_types parametrization covers apply_pair_exclusion as intended.
🧹 Nitpick comments (7)
deepmd/dpmodel/model/make_model.py (1)
316-322: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring update looks accurate; stale example nearby.
Matches the new DPA1 graph-native attention behavior (attention layers now included in graph eligibility). Note the unchanged
_call_common_graphexception message a few dozen lines below ("e.g. dpa1 attn_layer=0") is now a narrower example than what this docstring describes — consider updating that message text for consistency in a follow-up.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/model/make_model.py` around lines 316 - 322, The exception message in _call_common_graph is now too narrow compared with the updated graph-native attention behavior described in the nearby docstring. Update the message text in _call_common_graph so it reflects the broader DPA1 attention-layer graph eligibility instead of only referencing the old “e.g. dpa1 attn_layer=0” example, keeping the wording consistent with the behavior documented in make_model.py.source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py (1)
166-179: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest comment overstates what is actually verified.
The comment says this block will "verify excluded pairs contribute sw == 0" and "check ... call_graph sw channel", but
call_graphonly returns(grrg, rot_mat)— it has noswoutput — and the actual assertions only check for NaN/Inf, not the claimed masking behavior. The earlierout[4]vsref[4]comparison (lines 162-164) already indirectly validates exclusion parity forswvia the dense reference, so this block is largely redundant and its comment is misleading about intent/coverage. Either remove the stale comment or replace it with an assertion that actually validates zeroed contributions from excluded pairs (e.g., inspect the block'sedge_mask/sw_eviase_atten.call_graphdirectly).♻️ Suggested comment fix (minimal)
- if exclude_types: - # verify excluded pairs contribute sw == 0 in the dense reference - # (atype=[0,1,0,1] -> pairs (0,1) and (1,0) should be masked) - # sw shape: (nf, nloc, nnei, 1); just check the graph output is also 0 - # for excluded-pair edges by checking call_graph sw channel + if exclude_types: + # additional sanity check on the raw call_graph output (no sw + # channel here; exclusion parity for sw is already verified via + # out[4] vs ref[4] above). graph = from_dense_quartet(ext_coord, nlist, mapping, compact=False) atype_local = self.atype.reshape(-1) - grrg_g, rot_mat_g = dd.call_graph( + grrg_g, _rot_mat_g = dd.call_graph( graph, atype_local, type_embedding=dd.type_embedding.call() ) # no nan/inf in output with exclusions applied assert not np.any(np.isnan(grrg_g)) assert not np.any(np.isinf(grrg_g))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py` around lines 166 - 179, The comment in test_dpa1_call_graph_descriptor is misleading because this block does not verify excluded-pair sw masking; call_graph only returns grrg and rot_mat, and the current assertions only check NaN/Inf. Update the test by either removing/rephrasing the stale comment to match the actual coverage, or add a real assertion for zeroed excluded-pair contributions by checking the relevant sw/edge-mask path through se_atten.call_graph or the returned graph data. The earlier out[4] vs ref[4] comparison already covers sw parity, so keep this block focused on what it truly validates.source/tests/common/dpmodel/test_neighbor_graph_builder.py (1)
419-427: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant
import unittest.
unittestis already imported at the top of this file; the local re-import inside theexceptblock is unnecessary.🧹 Proposed cleanup
`@classmethod` def setUpClass(cls) -> None: try: import ase # noqa: F401 except ImportError as e: - import unittest - raise unittest.SkipTest("ase not installed") from e🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_neighbor_graph_builder.py` around lines 419 - 427, Remove the redundant local import inside test_neighbor_graph_builder’s setUpClass method: the file already imports unittest, so keep the ImportError handling but drop the inner import and use the existing unittest.SkipTest reference when ase is missing.Source: Linters/SAST tools
deepmd/dpmodel/utils/neighbor_graph/graph.py (1)
192-194: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring overstates what
compact=Truereplaces.The parameter doc says
edge_index,edge_vec,angle_index,angle_maskare all replaced whencompact=True. In practiceangle_index/angle_maskare never touched by the compact branch — the function only reaches the compaction step after confirming both areNone(otherwise it raisesNotImplementedError). Listing them as "replaced" could mislead a future implementer extending angle-compaction support into thinking this path already handles it.📝 Suggested doc fix
graph - The neighbor graph; only ``edge_mask`` (and, if ``compact=True``, - ``edge_index``, ``edge_vec``, ``angle_index``, ``angle_mask``) are - replaced. + The neighbor graph; only ``edge_mask`` (and, if ``compact=True``, + ``edge_index`` and ``edge_vec``) are replaced. ``angle_index`` / + ``angle_mask`` are never touched — compaction is rejected outright + when either is present (see the ``compact`` behavior below).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/utils/neighbor_graph/graph.py` around lines 192 - 194, The docstring for the neighbor graph parameter overstates the effect of compact=True by implying that angle_index and angle_mask are also replaced. Update the documentation in graph.py to say compact mode only compacts edge_index and edge_vec (along with edge_mask), and make it clear that angle_index and angle_mask are not handled by this branch because the code path only proceeds when they are None.deepmd/dpmodel/utils/neighbor_graph/ase_builder.py (1)
154-163: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winPin device explicitly when converting
atypeforapply_pair_exclusion.
xp = array_api_compat.array_namespace(coord)followed byxp.asarray(atype)doesn't pin a device, unlike the analogouspair_exclwiring innv_graph_builder.pyandvesin_graph_builder.py, which both usetorch.as_tensor(atype, device=<coord's device>). Ifatypeisn't already a tensor on the same device ascoord(e.g. a CPU/numpyatypepaired with a CUDAcoord),xp.asarraywill silently produce a CPU tensor, which will then device-mismatch againstgraph.edge_index/edge_maskinsideapply_pair_exclusion.🔧 Suggested fix
if pair_excl is not None: import array_api_compat xp = array_api_compat.array_namespace(coord) - atype_flat = xp.reshape(xp.asarray(atype), (-1,)) + dev = array_api_compat.device(coord) + atype_flat = xp.reshape(xp.asarray(atype, device=dev), (-1,)) graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) return graph🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/utils/neighbor_graph/ase_builder.py` around lines 154 - 163, The atype conversion in the ASE neighbor graph path is not explicitly pinned to coord’s device, so apply_pair_exclusion can receive tensors on the wrong device. Update the ase_builder flow that builds graph and handles pair_excl to convert atype the same way as the nv_graph_builder and vesin_graph_builder paths: derive the device from coord and create atype on that device before flattening and passing it into apply_pair_exclusion. This keeps the device consistent with graph.edge_index and graph.edge_mask.source/tests/common/dpmodel/test_graph_atomic_parity.py (1)
318-344: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the unused model scaffolding.
amis never referenced here, soDescrptDPA1,InvarFitting, andDPAtomicModelcan be removed from this test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_graph_atomic_parity.py` around lines 318 - 344, The test builds unused model scaffolding that is never referenced, so remove the dead setup from test_apply_pair_exclusion_idempotent. Eliminate the DescrptDPA1, InvarFitting, and DPAtomicModel construction (including the am variable) and keep only the inputs actually needed for extend_input_and_build_neighbor_list, from_dense_quartet, and apply_pair_exclusion. Make sure the test still covers both the empty and non-empty pair_exclude_types branches.Sources: Coding guidelines, Linters/SAST tools
source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc (1)
101-159: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftTest coverage gap: LAMMPS
InputNlistingestion route not exercised for pair-exclusion.
check_against_ref/allTYPED_TESTs call the 6-argdp.compute(ener, force, virial, coord, atype, box), which routes toDeepPotPTExpt's standalone (no-nlist,build_nlist-based)compute()overload. The LAMMPS-styleInputNlistoverload — the actual pair-style ingestion seam, which cachesedge_index_tensor/firstneigh_tensoratago==0and recomputes geometry viacompactEdgeTensorsevery step before callingapplyPairExclusion/applyPairExclusionNlist— is never invoked here. A bug isolated to that branch's node/edge tensor construction (e.g. themulti_rank ? nall_real : nlocnode-count selection feedingapplyPairExclusion) wouldn't be caught by this suite.Consider adding a case that drives the
InputNlistoverload (mirroring the pattern intest_deeppot_dpa1_graph_ptexpt.cc) withpair_exclude_typesset, so both C++ ingestion entry points are validated against the Python reference.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc` around lines 101 - 159, Add coverage for the LAMMPS-style InputNlist ingestion path in this pair-exclusion test, because the current check_against_ref and TYPED_TESTs only exercise the 6-arg DeepPot::compute route. Introduce a test that calls the InputNlist compute overload on DeepPotPTExpt, using pair_exclude_types and matching the pattern used in test_deeppot_dpa1_graph_ptexpt.cc, so the edge/node tensor caching and applyPairExclusion/applyPairExclusionNlist branch are validated against the Python reference.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt_expt/entrypoints/main.py`:
- Around line 571-576: Update the stale inline comment in main by removing the
“no type exclusion” restriction so it matches the current graph-eligibility
behavior. Keep the note aligned with `_model_uses_graph_lower` in training.py
and the nearby ValueError message: describe graph lower as opt-in for
graph-eligible models (dpa1 with concat tebd, attention layers, and supported
exclude_types) and preserve the rest of the fail-fast/per-atom-virial
explanation.
In `@doc/model/train-se-atten.md`:
- Around line 160-164: Update the pt_expt training doc sentence describing
graph-eligible descriptors so it no longer says descriptor-level exclude_types
disqualifies the carry-all neighbor-graph path. Use the surrounding
se_atten/neighbor_graph_method explanation to state that mixed-type descriptors
with tebd_input_mode "concat" and no descriptor-level compression remain
graph-eligible, while exclude_types is not a blocking condition anymore. Keep
the dense-vs-graph parity note tied to smooth_type_embedding and attn_layer, but
make the eligibility rule consistent with the current behavior exercised by
test_exclude_types_graph_eligible_and_parity and dd.uses_graph_lower().
---
Outside diff comments:
In `@source/tests/pt_expt/descriptor/test_dpa1.py`:
- Around line 117-147: test_exportable is missing the graph mapping input, so it
still exercises the dense fallback instead of the graph-native exclusion path.
Update the export setup in test_exportable to pass mapping into
torch.export.export alongside dd0 and the existing inputs, using the test
fixture’s mapping source so the exclude_types parametrization covers
apply_pair_exclusion as intended.
---
Nitpick comments:
In `@deepmd/dpmodel/model/make_model.py`:
- Around line 316-322: The exception message in _call_common_graph is now too
narrow compared with the updated graph-native attention behavior described in
the nearby docstring. Update the message text in _call_common_graph so it
reflects the broader DPA1 attention-layer graph eligibility instead of only
referencing the old “e.g. dpa1 attn_layer=0” example, keeping the wording
consistent with the behavior documented in make_model.py.
In `@deepmd/dpmodel/utils/neighbor_graph/ase_builder.py`:
- Around line 154-163: The atype conversion in the ASE neighbor graph path is
not explicitly pinned to coord’s device, so apply_pair_exclusion can receive
tensors on the wrong device. Update the ase_builder flow that builds graph and
handles pair_excl to convert atype the same way as the nv_graph_builder and
vesin_graph_builder paths: derive the device from coord and create atype on that
device before flattening and passing it into apply_pair_exclusion. This keeps
the device consistent with graph.edge_index and graph.edge_mask.
In `@deepmd/dpmodel/utils/neighbor_graph/graph.py`:
- Around line 192-194: The docstring for the neighbor graph parameter overstates
the effect of compact=True by implying that angle_index and angle_mask are also
replaced. Update the documentation in graph.py to say compact mode only compacts
edge_index and edge_vec (along with edge_mask), and make it clear that
angle_index and angle_mask are not handled by this branch because the code path
only proceeds when they are None.
In `@source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc`:
- Around line 101-159: Add coverage for the LAMMPS-style InputNlist ingestion
path in this pair-exclusion test, because the current check_against_ref and
TYPED_TESTs only exercise the 6-arg DeepPot::compute route. Introduce a test
that calls the InputNlist compute overload on DeepPotPTExpt, using
pair_exclude_types and matching the pattern used in
test_deeppot_dpa1_graph_ptexpt.cc, so the edge/node tensor caching and
applyPairExclusion/applyPairExclusionNlist branch are validated against the
Python reference.
In `@source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py`:
- Around line 166-179: The comment in test_dpa1_call_graph_descriptor is
misleading because this block does not verify excluded-pair sw masking;
call_graph only returns grrg and rot_mat, and the current assertions only check
NaN/Inf. Update the test by either removing/rephrasing the stale comment to
match the actual coverage, or add a real assertion for zeroed excluded-pair
contributions by checking the relevant sw/edge-mask path through
se_atten.call_graph or the returned graph data. The earlier out[4] vs ref[4]
comparison already covers sw parity, so keep this block focused on what it truly
validates.
In `@source/tests/common/dpmodel/test_graph_atomic_parity.py`:
- Around line 318-344: The test builds unused model scaffolding that is never
referenced, so remove the dead setup from test_apply_pair_exclusion_idempotent.
Eliminate the DescrptDPA1, InvarFitting, and DPAtomicModel construction
(including the am variable) and keep only the inputs actually needed for
extend_input_and_build_neighbor_list, from_dense_quartet, and
apply_pair_exclusion. Make sure the test still covers both the empty and
non-empty pair_exclude_types branches.
In `@source/tests/common/dpmodel/test_neighbor_graph_builder.py`:
- Around line 419-427: Remove the redundant local import inside
test_neighbor_graph_builder’s setUpClass method: the file already imports
unittest, so keep the ImportError handling but drop the inner import and use the
existing unittest.SkipTest reference when ase is missing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: f66e001c-837b-4b14-a6b8-84d867cc8a56
📒 Files selected for processing (48)
deepmd/dpmodel/array_api.pydeepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/model/spin_model.pydeepmd/dpmodel/utils/__init__.pydeepmd/dpmodel/utils/default_neighbor_list.pydeepmd/dpmodel/utils/neighbor_graph/__init__.pydeepmd/dpmodel/utils/neighbor_graph/ase_builder.pydeepmd/dpmodel/utils/neighbor_graph/builder.pydeepmd/dpmodel/utils/neighbor_graph/env.pydeepmd/dpmodel/utils/neighbor_graph/graph.pydeepmd/dpmodel/utils/neighbor_graph/pairs.pydeepmd/dpmodel/utils/neighbor_graph/segment.pydeepmd/dpmodel/utils/neighbor_list.pydeepmd/dpmodel/utils/nlist.pydeepmd/pt/utils/nv_nlist.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/nv_graph_builder.pydeepmd/pt_expt/utils/serialization.pydeepmd/pt_expt/utils/vesin_graph_builder.pydeepmd/pt_expt/utils/vesin_neighbor_list.pydoc/model/train-se-atten.mdsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/include/commonPT.hsource/api_cc/src/DeepPotPTExpt.ccsource/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.ccsource/install/test_cc_local.shsource/tests/common/dpmodel/test_apply_pair_exclusion.pysource/tests/common/dpmodel/test_apply_pair_exclusion_nlist.pysource/tests/common/dpmodel/test_center_edge_pairs.pysource/tests/common/dpmodel/test_dpa1_call_graph_block.pysource/tests/common/dpmodel/test_dpa1_call_graph_descriptor.pysource/tests/common/dpmodel/test_dpa1_graph_attention_parity.pysource/tests/common/dpmodel/test_graph_atomic_parity.pysource/tests/common/dpmodel/test_neighbor_graph_builder.pysource/tests/common/dpmodel/test_segment_softmax.pysource/tests/common/dpmodel/test_spin_model_legacy_routing.pysource/tests/infer/gen_dpa1_pairexcl.pysource/tests/pt_expt/descriptor/test_dpa1.pysource/tests/pt_expt/infer/test_graph_deepeval.pysource/tests/pt_expt/model/test_dpa1_graph_lower.pysource/tests/pt_expt/model/test_linear_model.pysource/tests/pt_expt/utils/test_graph_pt2_metadata.pysource/tests/pt_expt/utils/test_neighbor_list.pysource/tests/pt_expt/utils/test_vesin_graph_builder.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5733 +/- ##
==========================================
- Coverage 79.62% 79.52% -0.11%
==========================================
Files 1014 1015 +1
Lines 115533 115818 +285
Branches 4276 4291 +15
==========================================
+ Hits 91995 92102 +107
- Misses 21994 22164 +170
- Partials 1544 1552 +8 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
58f771f to
9f2c31a
Compare
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
d5b5032 to
9c96566
Compare
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
…ct=True) when angle fields are present
…exclusion; drop eligibility gate
…layer and type_one_side Add @pytest.mark.parametrize for attn_layer in [0, 2] and type_one_side in [False, True] to test_exclude_types_graph_parity. Also adds the missing parity assertion (graph vs dense at rtol=atol=1e-12, non-binding sel). Uses smooth_type_embedding=False to avoid the known by-design softmax denominator divergence in the dense smooth path.
Descriptor-level exclude_types is now graph-eligible (fully supported via apply_pair_exclusion). Remove 'no exclude_types' from four docstrings/error messages that list graph eligibility conditions. The gate condition was removed in the NeighborGraph implementation; only tebd_input_mode='concat' restriction remains. - deepmd/pt_expt/entrypoints/main.py: freeze_model docstring (~502) + ValueError message (~589) - deepmd/dpmodel/model/make_model.py: forward docstring (~317) - deepmd/pt_expt/train/training.py: _model_uses_graph_lower docstring (~591)
build_neighbor_graph, build_neighbor_graph_ase, build_neighbor_graph_vesin, build_neighbor_graph_nv all gain optional keyword-only pair_excl=None and compact=False; default path = geometric search then apply_pair_exclusion. _call_common_graph in pt_expt make_model wires atomic_model.pair_excl to every builder call so model-level pair_exclude_types is applied at build time (the atomic-model seam backstop stays as idempotent identity). Oracle tests assert set-equality of the valid-edge set between builder(pair_excl=X) and builder() + separate apply_pair_exclusion(X), for dense (2 id + 3 oracle cases) and ase (2 cases); vesin gets 4 new tests (2 identity, 2 oracle, parametrized over periodic).
…n for array_api_strict compat
…_neighbor_list/strategies (A4) Extract the inline pair-exclusion from base_atomic_model.forward_common_atomic into apply_pair_exclusion_nlist(nlist, atype_ext, pair_excl) in nlist.py. The seam is refactored to call the named helper (idempotent backstop remains). Add pair_excl=None to: - build_neighbor_list (dpmodel, nlist.py) - DefaultNeighborList.build - VesinNeighborList.build (pt_expt) - NvNeighborList.build (pt; CUDA-only, API parity) - NeighborList base class signature 12 new unit tests covering: None/empty identity, excluded pairs -> -1, -1 slot preservation, ghost-atom types, idempotence, torch namespace smoke, build_neighbor_list oracle equivalence, DefaultNeighborList oracle, VesinNeighborList oracle. NvNeighborList CUDA-only (not validated locally).
…ce B) The dpa1 forward computes its env matrix through the NeighborGraph (from_dense_quartet -> edge_env_mat); the input stat used the dense EnvMat. Route the dpa1 block's input stat through the SAME graph path so stat and forward share one env-matrix implementation, with both exclusions folded in exactly as the forward does (model-level via the pre-excluded nlist from Piece A; descriptor-level via the emask mask). BIT-IDENTICAL to the dense path, so stored davg/dstd are unchanged: from_dense_quartet(compact=False) reuses the same neighbor set + padding (row-major (frame,center,slot) edges), edge_env_mat mirrors EnvMat.call, and the (E,4) output reshapes 1:1 to the dense (nf,nloc,nsel,4) tensor. Opt-in via EnvMatStatSe(use_graph=True); se_e2_a/se_r are untouched. pt_expt inherits it via autowrap; legacy pt stays dense (bit-identical, so cross-backend parity holds). Tests: graph==dense stat bit-identical (1e-15) for se_e2_a and the dpa1 block, under no exclusion, model-level pair_exclude, and descriptor-level exclude.
_graph_env_mat (input stat) duplicated the dense-quartet -> (graph, atype_local) setup from DescrptDPA1._call_graph_adapter almost verbatim (coord reshape, mapping-None identity, from_dense_quartet compact=False, xp_take_first_n local atype). Extract it as neighbor_graph.graph_from_dense_ quartet(coord_ext, atype_ext, nlist, mapping) -> (graph, atype_local) and route both call sites through it. Pure extraction, bit-identical (dpa1 adapter + stat parity tests unchanged).
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
24cdad6 to
06e6ede
Compare
for more information, see https://pre-commit.ci
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Thanks for pushing this toward a build-seam implementation. I agree with the direction, but I don't think this is safe to merge yet: several lower/precomputed-entry paths still bypass the new pair-exclusion seam, and the input-stat cache can be reused across different model-level pair_exclude_types.
Requesting changes so the exclusion behavior is closed over all external entry points before merge. CI being green is expected here because the missing paths are mostly with-comm / direct edge / precomputed-neighbor cases that the current parity tests don't exercise.
Reviewed by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5
| // never re-applies it; this is the single application site on the C++ | ||
| // dense route. | ||
| const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( | ||
| firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); |
There was a problem hiding this comment.
This only covers the non-comm dense route. The use_with_comm dense branch above still calls run_model_with_comm(..., firstneigh_tensor, ...) directly, so a .pt2 with pair_exclude_types will still include excluded pairs in multi-rank / with-comm inference. Please apply applyPairExclusionNlist before the with-comm call too, and add a with-comm regression test.
There was a problem hiding this comment.
Fixed in ff28370: applyPairExclusionNlist is now applied on firstneigh_tensor before run_model_with_comm, mirroring the single-rank dense site. Added a with-comm regression test on a message-passing DPA3 model (test_pair_deepmd_mpi_dpa3_pairexcl_matches_single_rank) — MP(-n 2) ≡ SP(-n 1) AND excluded ≠ no-exclusion baseline — validated on a clean GPU box. Full path-by-path matrix + rationale in the summary (cell 3).
| @@ -1238,13 +1268,26 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, | |||
| edge_tensors.edge_index_ext, edge_tensors.edge_mask, | |||
There was a problem hiding this comment.
The lower_input_kind == "edge_vec" paths still pass the original edge_tensors.edge_mask into run_model_edges* without applying pair_exclude_table_. This affects both this non-comm path and the with-comm edge path above. Since the exported lower no longer re-applies model-level pair_exclude_types, the edge schema needs an applyPairExclusion-equivalent mask at this ingestion seam as well. Please be careful about the type space: fold_to_local=false needs extended atypes, while folded graph/edge inputs need local atypes.
There was a problem hiding this comment.
Out of scope — the edge lower is a pt-backend export (SeZM/DPA4 via deepmd/pt/entrypoints/freeze_pt2.py, lower_input_kind="edge_vec"); pt_expt serialization only emits "graph"/"nlist". This PR does not touch the pt backend (its whole pt footprint is nv_nlist.py, +28), and the pt backend still applies exclusion at consume time (pt/model/atomic_model/base_atomic_model.py:363-366), so the edge .pt2 bakes exclusion into the exported graph — the C++ edge consumers correctly do NOT re-apply it (that would double-apply). So exclusion isn't dropped here. Details in the summary ("Out of scope — the edge lower").
| selection is a pure performance choice and results are unchanged. | ||
| """ | ||
| method = self._neighbor_graph_method | ||
| # Model-level ``pair_exclude_types`` is a graph-BUILD transform |
There was a problem hiding this comment.
This wires model-level exclusion into the graph/nlist builders, but the direct edge fast path earlier in _prepare_lower_inputs still bypasses it: self._nlist_builder.build(..., return_mode="edges") returns an edge schema that is passed straight to the edge lower. vesin_neighbor_list.py also rejects pair_excl for return_mode="edges", so the caller needs to post-filter the returned edge schema / mask with the same graph-style exclusion. Please add a DeepEval lower_input_kind="edge_vec" + vesin/nv + pair_exclude_types test.
There was a problem hiding this comment.
Same as the C++ edge path: this consumes a pt-backend edge .pt2 (SeZM/DPA4) whose exported graph already bakes in exclusion (the pt backend still applies it at consume time), so the return_mode="edges" fast path correctly does not post-filter. Out of scope for this dpmodel-derived-backend refactor. See the summary.
| # nlist. Excluded pairs then behave exactly like empty slots | ||
| # (env_mat 0, still counted) -- identical to descriptor-level | ||
| # exclude_types, replacing the previous accumulation-deselect. | ||
| pair_excl=pair_excl, |
There was a problem hiding this comment.
Now that input statistics depend on model-level pair_exclude_types through the pre-excluded nlist, the stat cache key also needs to include a canonicalized exclusion set. EnvMatStatSe.get_hash() still hashes descriptor shape/cutoff/sel/etc. but not pair_exclude_types, so changing only the exclusion list can silently reuse stale stats.
There was a problem hiding this comment.
Valid, but pre-existing and deferred: stats depended on pair_exclude_types before this PR too (via the accumulation-deselect this PR replaced with a build-time nlist mask), and get_hash() never included it. It's also not cleanly fixable here — the model-level exclusion isn't visible at get_hash time (it arrives in sample data; path / get_hash() is computed before sampling). Tracking as a separate fix so this PR stays scoped to the regressions it introduced. See the summary ("Deferred — stat-cache hash").
|
One more blocker I could not place as an inline review comment because the file is not touched in this PR diff:
Please either apply the model's pair exclusion to the converted Reviewed by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5 |
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Model-level pair_exclude_types is silently dropped on the C++ multi-rank (message-passing) dense path and the edge-schema paths
This PR moves model-level pair_exclude_types from a consume-time transform to a build-time (ingestion-seam) transform: BaseAtomicModel.forward_common_atomic / forward_common_atomic_graph no longer apply it, so the exported .pt2 lower no longer carries it and the C++ ingestion seam becomes the only application site.
But the seam was added to only two of the run paths in DeepPotPTExpt::compute:
- ✅ single-rank graph →
run_model_graph(applyPairExclusiononedge_mask) - ✅ single-rank dense →
run_model(applyPairExclusionNliston the nlist)
The following paths pass the raw tensors and therefore silently ignore model-level exclusion:
- ❌
run_model_with_comm— multi-rank dense (DeepPotPTExpt.cc:842).use_with_comm = has_comm_artifact_ && multi_rank(:521), i.e. exactly the message-passing (DPA2/DPA3) path undermpirun. - ❌
run_model_edges/run_model_edges_with_comm— edge schema, e.g. SeZM (:822/:835/:851/:1265).
Consequences
- A message-passing model with
pair_exclude_typesgives multi-rank ≠ single-rank, and it is a regression vs. before this PR (exclusion used to be baked into the.pt2, covering every path). - Edge-schema models lose exclusion on every rank.
- No fail-fast — the output is silently wrong.
dpa1 (this PR's focus) is unaffected: it has no comm artifact, so multi-rank still routes through run_model. The gap is latent for the model types this PR doesn't target, but the shared base-class change degrades them with no guard.
Why it's easy to miss — the narrative contradicts the code
- The PR description says the atomic-model seam "stays as idempotent backstop" and "Multi-rank (mpirun) exclusion shares the same seam"; neither holds now.
vesin_neighbor_list.pydocstring andpt_expt/model/make_model.py(~L472) still describe an "idempotent backstop", contradicting the "NOT applied here" comments inbase_atomic_model.py.
Requested fix (keeping decision #18's single-build-seam design)
Apply the transform on every C++ run path — applyPairExclusionNlist on the run_model_with_comm nlist, and applyPairExclusion on the edge_mask for the edge / edge-with-comm paths — and add a fail-fast for any pair_exclude_types × path combination that genuinely cannot be covered. Please also reconcile the stale "backstop" comments/docstrings so the single-owner contract reads consistently.
| // #18/A4): the exported dense lower consumes a pre-excluded nlist and | ||
| // never re-applies it; this is the single application site on the C++ | ||
| // dense route. | ||
| const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( |
There was a problem hiding this comment.
Single-rank dense applies exclusion here, but its siblings do not: run_model_with_comm (L842, multi-rank dense) and run_model_edges / run_model_edges_with_comm (L822 / L835 / L851 / L1265, edge schema) pass the raw tensors. Since the exported .pt2 no longer applies model-level exclusion (removed from forward_common_atomic), those paths silently drop it, so multi-rank ≠ single-rank for message-passing models and edge-schema models lose it entirely. Please thread applyPairExclusionNlist / applyPairExclusion through them too (or fail-fast).
There was a problem hiding this comment.
Addressed in ff28370. The dense multi-rank path (run_model_with_comm) now applies applyPairExclusionNlist — verified on a message-passing DPA3 with-comm model (MP≡SP + active-vs-baseline, clean-env). The edge-schema paths are intentionally left alone: they're a pt-backend lower (SeZM/DPA4) that still bakes exclusion into the exported graph, so re-applying in C++ would double-apply. Stale "backstop" docstrings reconciled. Full SP/MP × dense/graph matrix + the fail-fast reasoning for cell 6 in the summary.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
C++ hot path uploads the constant pair-exclusion table to the device every step
applyPairExclusion (commonPT.h:547-551) and applyPairExclusionNlist (:602-606) rebuild the keep table from the std::vector and copy it to the device on every compute() call (i.e. every MD step):
const auto table =
torch::from_blob(const_cast<int*>(type_mask_table.data()),
{static_cast<std::int64_t>(type_mask_table.size())},
torch::TensorOptions().dtype(torch::kInt32))
.clone()
.to(device);pair_exclude_table_ is a compile-time constant fixed in init(), and the device is fully determined at init (gpu_id / gpu_enabled), so this is a per-step CPU clone + H2D upload (plus a stream op / possible sync on CUDA) of a value that never changes. It is guarded by the empty-table early-return, so only exclusion-enabled models pay it — but for those it is pure per-step waste on the hot path.
Proposed fix — upload once at init, pass the device tensor in
Store the table as a device-resident torch::Tensor built once, and let the two helpers consume it directly (an undefined tensor => identity, mirroring the current empty-vector early-exit). buildPairExcludeTable stays a pure, unit-testable vector builder.
1. Member (DeepPotPTExpt.h) — change the type:
// Device-resident (ntypes+1)^2 keep table, uploaded once in init.
// Undefined tensor => no model-level exclusion.
torch::Tensor pair_exclude_table_;2. One-time upload in init (device known here; same construction as compute() at :463-466):
std::vector<int> tbl = deepmd::buildPairExcludeTable(ntypes, pair_exclude_types);
if (!tbl.empty()) {
torch::Device device(torch::kCUDA, gpu_id);
if (!gpu_enabled) {
device = torch::Device(torch::kCPU);
}
pair_exclude_table_ =
torch::from_blob(tbl.data(),
{static_cast<std::int64_t>(tbl.size())}, torch::kInt32)
.clone()
.to(device);
}3. Helpers take the prebuilt tensor and drop the per-call from_blob/clone/to:
inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index,
const torch::Tensor& edge_mask,
const torch::Tensor& atype,
const torch::Tensor& type_mask_table,
const int ntypes) {
if (!type_mask_table.defined()) {
return edge_mask; // identity
}
const auto src = edge_index.index({0});
const auto dst = edge_index.index({1});
const auto type_ij =
atype.index_select(0, dst) * (ntypes + 1) + atype.index_select(0, src);
const auto keep = type_mask_table.index_select(0, type_ij).to(torch::kBool);
return torch::logical_and(edge_mask, keep);
}applyPairExclusionNlist mirrors this (keep = type_mask_table.index_select(0, type_ij.reshape({-1})).reshape({nf, nloc, nnei})). The table already lives on the model device (== edge_mask / nlist device), so index_select needs no per-call transfer.
This keeps the single-owner design intact, removes the per-step allocation + H2D entirely, and is self-contained: the helpers are called only inside DeepPotPTExpt::compute (the gtest drives them through the public API), so no other call site changes.
There was a problem hiding this comment.
The one canonical exclusion transform is a sound design, but moving it to the build seam and dropping the forward_common_atomic* backstop makes exclusion fail-open: those methods now document a "nlist/graph is already pre-excluded" contract, so any ingestion entry point that skips the build seam silently includes excluded pairs instead of erroring. For a security-relevant feature (exclusion forces zero interaction between types), silent inclusion is the dangerous failure mode — and the leaks are being found one at a time (jax_md call_lower, DeepEval return_mode="edges", the C++ with-comm / edge-schema paths, the env_mat_stat hash key already flagged by njzjz-bot / OutisLi), which suggests there's no enumerated inventory of entry points.
Requesting changes on that architectural point: could we either keep a cheap fail-safe guard (an eager/debug assertion that inputs entering forward_common_atomic* are exclusion-applied), or enumerate every ingestion path in this PR with one regression test each, so coverage is complete-by-construction? The contract boundary is also currently untested (all exclusion tests go through the pre-excluding build seam), and the PR description's "stays as idempotent backstop" wording is now inaccurate at HEAD ("never re-applied here") — worth correcting so the blast radius reads right.
…) path Address the pair_exclude_types review on deepmodeling#5733 (njzjz-bot #1, OutisLi). Decision deepmodeling#18/A4 moved model-level pair_exclude_types to a build-time C++ ingestion seam, but the seam was wired into only the single-rank dense (applyPairExclusionNlist) and graph (applyPairExclusion) routes. The multi-rank dense path (run_model_with_comm, DeepPotPTExpt.cc) passed the raw nlist, so a message-passing .pt2 with pair_exclude_types silently included excluded pairs multi-rank (multi-rank != single-rank). Apply the SAME seam before run_model_with_comm, mirroring the single-rank dense site exactly (the cross-rank ghost exchange happens inside run_model_with_comm and does not change the nlist's meaning, so per-rank pre-exclusion is correct). Scope: the edge-schema lowers (run_model_edges*) are intentionally NOT touched -- the edge lower is a pt-backend export (SeZM/DPA4 via pt/entrypoints/freeze_pt2.py) and the pt backend still applies exclusion at consume time, so the edge .pt2 bakes exclusion into the graph. The MP message-passing GRAPH lower stays fail-fast (PR-G). Perf (OutisLi): buildPairExcludeTable's flat table was rebuilt and copied H2D on every compute() (MD) step. Upload it once in init as a device torch::Tensor member (pair_exclude_table_); the seam helpers now take the prebuilt tensor and index_select it directly (undefined tensor => identity, mirroring the old empty-vector early-exit). No per-step allocation or transfer. Docs: reconcile the stale "idempotent backstop" wording (vesin_neighbor_list, pt_expt/model/make_model) with the single-owner design -- both dpmodel seams (forward_common_atomic{,_graph}) now say "NOT applied here"; the build site is the sole owner. Test (cell 5): add a dpa1 MP-graph + pair_exclude_types parity test to test_lammps_dpa1_graph_pt2.py -- MP (-n 2) == SP (-n 1) on the excluded model, plus a cross-check that the excluded run differs from the same-weights no-exclusion baseline (so a silently-dropped exclusion on both ranks cannot pass trivially). Reuses the existing MPI harness with a pb= override. Verified locally (CPU): 8/8 single-rank pairexcl gtests pass after the perf refactor; all 6 tests in test_lammps_dpa1_graph_pt2.py pass (incl. the new multi-rank exclusion test).
Add deeppot_dpa3_pairexcl_mpi.pt2 (gen_dpa3.py, same weights as deeppot_dpa3_mpi.pt2 + pair_exclude_types=[[0,1]]) and an MP==SP + active-vs- baseline test exercising the run_model_with_comm pair-exclusion seam (cell 3). To be validated on a clean env; local build_cpu MPI is inconsistent.
for more information, see https://pre-commit.ci
Pair-exclusion coverage across the C++ run pathsThanks both — this was the right thing to flag. The build-seam refactor (decision #18/A4) made the C++ ingestion seam the single owner of model-level
Cell 6 is not a silent drop — it hard-throws ( Fixed — cells 3 & 5
Perf (OutisLi)The per-step table rebuild + H2D upload is removed: Out of scope — the edge lower (njzjz-bot #2, #3)The edge-schema lower ( Deferred — stat-cache hash (njzjz-bot #4)Valid, but pre-existing: input statistics depended on model-level Also (OutisLi): stale "idempotent backstop" wording in
|
Address iProzd's architectural review on deepmodeling#5733: decision deepmodeling#18/A4 moved pair_exclude_types to the build seam and dropped the consume-time backstop in forward_common_atomic{,_graph}, making exclusion fail-OPEN -- any ingestion path that skips the build seam would silently INCLUDE excluded pairs (the dangerous direction for an exclusion feature). Add _assert_nlist_pair_excluded / _assert_graph_pair_excluded: when pair_excl is set, verify no real neighbour/active edge carries an excluded type pair (a leak = the build seam was skipped) and raise a clear AssertionError otherwise. numpy-eager ONLY -- gated on array_api_compat.is_numpy_array, so it is a no-op under torch.export / jax jit (data-dependent, can't be traced) and in compiled production; the exported-.pt2 / C++ ingestion paths are covered by their own ingestion-site regression tests. Verified pt_expt make_fx/export still traces (guard skipped for torch tensors). Rewrite the two seam-contract tests (test_graph_atomic_parity) to assert the guard REJECTS a non-excluded input (the previously-untested contract boundary), keeping the build-time positive control; fix test_dp_atomic_model test_self_consistency to feed a pre-excluded nlist (matches test_excl_consistency). 641 common/dpmodel tests pass.
Close the jax-md leak iProzd flagged: _eval_with_jax_md_neighbor built the lower nlist from the JAX-MD neighbor list without model-level pair_exclude_types, then called call_lower -- which no longer re-applies it (decision deepmodeling#18/A4) -- so excluded pairs were silently included (fail-open). Fold apply_pair_exclusion_nlist in at this ingestion seam, mirroring DeepEval's nlist path. Guarded on the model actually carrying pair_excl, so the mock-model jax-md tests are unaffected. Test (runs without jax_md via the DenseNeighbor path): an excluded (0,1) pair drops the only edge -> energy 0 vs 0.04 without exclusion.
…at-graph-pair-exclude
|
@iProzd — thanks, this is the right framing and I agree: dropping the consume-time backstop made exclusion fail-open, which is why the leaks were surfacing one at a time. I've taken both remedies you offered — a fail-safe guard and an enumerated inventory — plus the doc fix. 1. Eager fail-safe guard (
|
…deepmodeling#5747) ## Summary Adds NeighborGraph (graph-native lower) support for the dpa1 descriptor with `tebd_input_mode="strip"`, closing the last descriptor-level gap that forced strip-mode models (and `se_atten_v2`, which is strip-by-construction) onto the legacy dense path. The dense strip branch factorizes the per-neighbor feature as `gg = gg_s*gg_t + gg_s` — a radial-only geometric embedding times a type-pair strip embedding (optionally switch-smoothed). Because this has **no neighbor-axis coupling**, it maps to the graph path edge-for-edge. The change is: - **Kernel** (`DescrptBlockSeAtten`): a new per-edge helper `_graph_edge_gg_strip` (op-for-op mirror of the dense strip branch, including the `center*ntypes_pad + nei` nei-fastest two-side table layout and `int64` gather indices), selected by a `concat`/`strip` branch in `call_graph`. - **Routing** (`DescrptDPA1.uses_graph_lower`): admits `strip`, while keeping compressed descriptors and `exclude_types` on the dense path (they have no graph kernel here). `se_atten_v2` inherits this and becomes graph-eligible for free. - **pt_expt**: the two graph `make_fx` export tests are parametrized over `tebd_input_mode` to prove the strip kernel is fx-traceable. No new op, no attention/`segment_sum` change, no C++/serialization change. ## Scope This PR is deliberately **independent of deepmodeling#5733** (graph `exclude_types`). It does **not** change `exclude_types` eligibility — the `uses_graph_lower` `and not exclude_types` gate and the `call_graph` `exclude_types` raise are both kept. When both land, whichever merges second resolves a small (2–3 line) mechanical conflict at `uses_graph_lower` / the `call_graph` guard. ## Test plan - Block-level graph-vs-dense strip parity at `rtol=atol=1e-12` over `type_one_side × smooth` (attn=0) and `type_one_side` (attn=2, non-smooth). - Descriptor-level routed-`call` vs `_call_dense` parity over `type_one_side × smooth × attn_layer` (incl. attn=2 + smooth=True, bit-exact via the `static_nnei` adapter), plus a negative-contract gate test (compressed → dense, strip+`exclude_types` → dense). - `se_atten_v2` eligibility + graph-vs-dense parity (replaces the obsolete "strip stays dense" test). - pt_expt strip `make_fx` export; cross-backend consistency strip cases now route pt_expt through the graph adapter. Validated on **CPU** and on **GPU (Tesla T4, cuda:0)**: pt_expt dpa1 50 passed, consistency strip 22 passed + `se_atten_v2` 110 passed, dpmodel strip suites 46 passed. No tolerances relaxed, no tests skipped. ## Known limitations - **Compression** stays on the dense path by design (strip-only tabulation has no graph kernel); the gate excludes `self.compress`. - **`exclude_types`** stays dense (out of scope — owned by deepmodeling#5733). - **jax** graph lower remains energy-only (analytical force on the graph route is a separate follow-up). - The graph path's `segment_sum`→`index_add` is atomic/non-deterministic on CUDA (1–2 fp64 ULP), inherent to atomic scatter; GPU parity validated within tolerance. - Pre-existing (not introduced here): a softmax `RuntimeWarning` on the shared attention path (max over fully-masked segments), also present on the concat path. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Summary * **New Features** * Added graph-based execution support for the additional stripped type-embedding mode. * **Bug Fixes** * Updated graph routing eligibility: compressed descriptors and excluded-type configurations now reliably fall back to dense execution; graph routing is disabled when compression is enabled. * **Tests** * Added bit-exact parity tests between graph and dense paths for the new stripped mode (including routing eligibility checks). * Expanded FX graph export/trace coverage for both embedding modes. * Adjusted neighbor-list fallback validation with model-specific tolerance handling and added a new smooth variant. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
…xclude # Conflicts: # deepmd/dpmodel/descriptor/dpa1.py # deepmd/jax/model/hlo.py # source/tests/common/dpmodel/test_dpa1_call_graph_block.py # source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py # source/tests/pt_expt/descriptor/test_dpa1.py
…no recompile The pair-exclusion DPA3 MPI fixture was a second full inductor compile (~2 with- comm compiles, expensive). Model-level pair_exclude_types is applied at the C++ ingestion seam (from metadata.json) and the Python DeepEval build seam (from the serialized model.json), NOT baked into the exported graph -- so the compiled AOTI artifact (incl. the nested with-comm .pt2) is byte-identical to deeppot_dpa3_mpi.pt2. Derive the pairexcl fixture by copying that archive and patching pair_exclude_types into BOTH JSON blobs (metadata.json for C++, model.json for Python) so the two paths agree. Verified via DeepEval: the derived model is bit-identical to a full recompile (E=2.5311 vs baseline 2.4916) and differs from the no-exclusion baseline. deeppot_dpa3_mpi.pt2 stays the exact no-exclusion baseline for the reference and active-vs-baseline tests.
Apply the metadata-patch fixture derivation (a95a958) to gen_dpa1_pairexcl.py and factor the logic into gen_common.derive_pair_exclude_pt2 (used by both gen_dpa3.py and gen_dpa1_pairexcl.py). gen_dpa1_pairexcl.py generated three inductor-compiled models (_none/_graph/_nlist). _graph has the same weights and lower_kind as the _none graph baseline and differs only in pair_exclude_types, so it is now DERIVED from _none by patching the archive -- no third compile. _nlist is a different lower_kind (different exported graph) so it stays a compile. 3 compiles -> 2. Verified via DeepEval: the derived _graph is bit-identical to a full recompile (E active vs the _none baseline). gen_dpa3.py refactored to use the shared helper.
OutisLi
dismissed
their stale review
Retracted by author; posted prematurely and partly out of scope.
Audit of the PR's added docstrings against the package numpydoc convention
(CLAUDE.md): public functions were already compliant; fix the private/helper
idiom mismatches so they match their files' sibling methods.
- base_atomic_model._assert_{nlist,graph}_pair_excluded: add Parameters + Raises
sections (siblings like _finalize_atomic_ret use full numpydoc).
- env_mat_stat._graph_env_mat: add Parameters + Returns (siblings iter/__call__
use numpydoc).
- jax2tf/make_model.model_call_from_call_lower: revert docstring to its original
one-liner (the function documents no other params); the pair_excl rationale is
already an inline comment at the application site, so the single-param prose
was redundant + inconsistent.
- gen_common.derive_pair_exclude_pt2: concise summary + typed Parameters, move
the decision/verification cross-refs into a Notes section (not floating prose).
No Google-style/See-Also/RTD-breaking issues found.
5 participants
Pair
exclude_typesas a canonical NeighborGraph transform (decision #18)Makes pair-type exclusion a single canonical transform applied once at the neighbor-list/graph build seam, and uses it to add descriptor-level
exclude_typessupport to the dpa1 graph path (removing that eligibility gate), consistently across dpmodel, pt_expt, jax, and the C++ inference path.What changed
apply_pair_exclusion(graph, atype, pair_excl, *, compact=False)indeepmd/dpmodel/utils/neighbor_graph/: ANDsPairExcludeMask.build_edge_exclude_maskintograph.edge_mask.compact=Falseis mask-only (shape-static, export/AOTI-safe);compact=Truedrops masked edges (eager-only; raises on angle-carrying graphs). Idempotent.base_atomic_model): model-levelpair_exclude_typesis applied at BUILD time, soforward_common_atomic{,_graph}no longer re-apply it — they consume a pre-excluded nlist/graph (single-owner, not a backstop). To keep that from being fail-open, an eager (numpy) fail-safe assertion (_assert_nlist_pair_excluded/_assert_graph_pair_excluded) rejects a non-excluded input at the seam with a clear error; it is a no-op undertorch.export/ jaxjitand in compiled production (the exported-.pt2/ C++ paths are covered by ingestion-site tests). See the ingestion-path inventory below.exclude_types: theNotImplementedErrorand theuses_graph_lower()exclude condition are removed; exclusion applied insideDescrptBlockSeAtten.call_graphbefore the segment sums. Graph-vs-dense parity at non-binding sel is exact (rtol=atol=1e-12, attn_layer 0 and 2, type_one_side both).build_neighbor_graphand the pt_expt graph builders (dense/ase/vesin/nv) gain optionalpair_excl/compactwith default post-search application;_call_common_graphpasses model-level excludes at build time; oracle set-equality tests per available builder.apply_pair_exclusion_nlist(nlist, atype_ext, pair_excl)extracted from the inline seam code;build_neighbor_list+ Vesin/Nv/Default strategies gainpair_excl;return_mode='edges'+pair_exclfails fast.buildPairExcludeTable/applyPairExclusion/applyPairExclusionNlistinsource/api_cc/include/commonPT.h, mirroring the Python transforms (same arg order/variable names, cross-referenced docs);pair_exclude_typesserialized into.pt2metadata.jsonand rebuilt once inDeepPotPTExpt::init(device-resident table, uploaded once — no per-step H2D). Exclusion is applied at the C++ ingestion seam — the single owner on every run path; the exported lower consumes a pre-excluded input and never re-applies it. New gtest (8 tests) vs Python DeepEval reference at 1e-10, plus multi-rank LAMMPS exclusion tests (below).apply_pair_exclusionuseslogical_and+ bool cast (array_api_strict rejectedbool*bool), caught by the jax/strict consistency rows now traversing the graph path.Ingestion-path inventory (exclusion coverage)
Because the seam is now the single owner (fail-open if skipped), here is every entry point that builds a nlist/graph reaching the exclusion-owning lower, and how each applies exclusion:
_call_common(dense)build_neighbor_list(pair_excl=)_call_common_graphbuild_neighbor_graph(pair_excl=)exclude_typesapply_pair_exclusionincall_graphapply_pair_exclusion_nlist/ vesinpair_excl_build_eval_graph(pair_excl=)(dense/ase/vesin/nv)build_neighbor_graph(pair_excl=)in stat/forwardbuild_neighbor_list(pair_excl=)inEnvMatStatSe.itercall_lowerapply_pair_exclusion_nlistat the seam (fixed here)test_dense_neighbor_applies_model_pair_exclusionapplyPairExclusionNlistapplyPairExclusionNlist(fixed here)applyPairExclusionapplyPairExclusionGuard:
forward_common_atomic{,_graph}additionally carry an eager fail-safe assertion (numpy only) that rejects a non-excluded input, so any future dpmodel/jax ingestion miss fails loudly instead of silently including excluded pairs.Known limitations
pair_exclpath has no local oracle test (CUDA-only); to be validated on a GPU box.pair_exclude_types— pre-existing (predates this PR) and tracked as a separate fix.build_edge_exclude_maskstill returns int32 (bool cast at call sites; follow-up).🤖 Generated with Claude Code
Spin routing
Spin models auto-inject
exclude_types(virtual/placeholder types) into their backbone descriptor; before this PR that condition accidentally kept spin on the dense path. With exclude_types now graph-eligible, spin backbones flipped onto the carry-all graph route, which (a) diverges from the sel-capped reference on sel-binding spin systems and (b) trips a torch-inductor scatter codegen assertion during spin.pt2export. Fixed explicitly:DescrptDPA1.disable_graph_lower()(not serialized; re-derived structurally) is set inSpinModel.__init__— the single choke point coveringget_spin_model,SpinModel.deserialize, and the pt_expt spin classes — plus a belt-and-bracesneighbor_graph_method="legacy"atSpinModel.call_common. Regression tests pin the routing and its serialize→deserialize survival; the full spin export suite (23) and spin checkpoint-interop suite (12) are green.Verification
Full pt_expt suite: 1196 passed / 39 skipped / 3 failed — the 3 failures (
test_dpa4_freeze_to_pt2,test_dpa4_deep_eval_*) are byte-identical on the base commit (pre-existing torch-inductor dpa4 export issue on this box, unrelated). dpmodel exclusion suites 69 passed; consistency dpa1 99 passed/63 skipped (incl. jax + array_api_strict exclude rows); C++Dpa1PairExclgtest 8/8.Summary by CodeRabbit
.pt2metadata and enforcement at inference ingestion.